import torch
from torch import nn



hazy_code, clear_code = codebook.chunk(chunks=2, dim=1)
hazy_resprestation, clear_resprestation = z.chunk(chunks=2, dim=1)
# (b,  h, w, 256)
hazy_resprestation = hazy_resprestation.permute(0, 2, 3, 1).contiguous()
# (b*h*w, 256)
hazy_resprestation = hazy_resprestation.view(-1, self.e_dim/2)
# (b*h*w, n_e每一个向量距离n_e个码元的距离)
hazy_d = self.dist(hazy_resprestation, hazy_code)

hazy_min_encoding_indices = torch.argmin(hazy_d, dim=1).unsqueeze(1)
hazy_min_encodings = torch.zeros(hazy_min_encoding_indices.shape[0], hazy_code.shape[0]).to(z)
hazy_min_encodings.scatter_(1, hazy_min_encoding_indices, 1)

# 完成去雾 使用在有雾上的索引矩阵✖️清晰码本=清晰
#(b*h*w, 256)
dehaze_resprestation_q = torch.matmul(hazy_min_encodings, clear_code)


